import time
from datetime import datetime
import json
import os
import uuid

import websockets
import numpy as np
import argparse
import ssl
from funasr.utils.postprocess_utils import rich_transcription_postprocess
from modelscope import pipeline
from whisper_live.transcriber import WhisperModel
from whisper_live.server import ServeClientFasterWhisper
from whisper_live.client import TranscriptionClient
from aiohttp import web
import ffmpeg

parser = argparse.ArgumentParser()
parser.add_argument("--host",
                    type=str,
                    default="0.0.0.0",
                    required=False,
                    help="host ip, localhost, 0.0.0.0")
parser.add_argument("--port",
                    type=int,
                    default=10095,
                    required=False,
                    help="grpc server port")
parser.add_argument("--httpport",
                    type=int,
                    default=10094,
                    required=False,
                    help="http server port")
parser.add_argument("--lang_model",
                    type=str,
                    default="damo/speech_campplus_five_lre_16k",
                    help="model from modelscope")
parser.add_argument("--lang_model_revision",
                    type=str,
                    default="v1.0.1",
                    help="")
parser.add_argument("--whisper_model",
                    type=str,
                    default="C:/Users/HP/.cache/huggingface/hub/faster_whisper",
                    help="model from modelscope")
parser.add_argument("--asr_model",
                    type=str,
                    default="damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch",
                    help="model from modelscope")
parser.add_argument("--asr_model_revision",
                    type=str,
                    default="v2.0.4",
                    help="")
parser.add_argument("--asr_model_online",
                    type=str,
                    default="damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-online",
                    help="model from modelscope")
parser.add_argument("--asr_model_online_revision",
                    type=str,
                    default="v2.0.4",
                    help="")
parser.add_argument("--vad_model",
                    type=str,
                    default="damo/speech_fsmn_vad_zh-cn-16k-common-pytorch",
                    help="model from modelscope")
parser.add_argument("--vad_model_revision",
                    type=str,
                    default="v2.0.4",
                    help="")
parser.add_argument("--punc_model",
                    type=str,
                    default="damo/punc_ct-transformer_zh-cn-common-vad_realtime-vocab272727",
                    help="model from modelscope")
parser.add_argument("--punc_model_revision",
                    type=str,
                    default="v2.0.4",
                    help="")
parser.add_argument("--ngpu",
                    type=int,
                    default=1,
                    help="0 for cpu, 1 for gpu")
parser.add_argument("--device",
                    type=str,
                    default="cuda",
                    help="cuda, cpu")
parser.add_argument("--ncpu",
                    type=int,
                    default=4,
                    help="cpu cores")
parser.add_argument("--certfile",
                    type=str,
                    default="../../ssl_key/server.crt",
                    required=False,
                    help="certfile for ssl")

parser.add_argument("--keyfile",
                    type=str,
                    default="../../ssl_key/server.key",
                    required=False,
                    help="keyfile for ssl")
parser.add_argument("--temp_dir",
                    type=str,
                    default="temp_dir",
                    required=False,
                    help="temp dir")
args = parser.parse_args()

websocket_users = set()


# 用于跟踪文件的分片上传状态
file_parts = {}

async def api_recognition_chunk(request):
    # 使用 multipart/form-data 解析请求内容
    post_data = await request.post()
    file_name = post_data["File-Name"]
    total_file_size = int(post_data["File-Size"])
    part_number = int(post_data["Part-Number"])
    chunk_size = int(post_data["Chunk-Size"])

    # 读取额外的表单字段
    lang = post_data.get("lang")
    print(f'lang={lang}')
    # 获取文件内容
    chunk_data = post_data["file"].file.read()  # 使用 .file.read() 获取文件数据
    part_path = os.path.join(f'{args.temp_dir}', f"{file_name}.part{part_number}")
    with open(part_path, "wb") as f:
        f.write(chunk_data)

    # 初始化或更新分片状态记录
    if file_name not in file_parts:
        total_parts = (total_file_size + chunk_size - 1) // chunk_size  # 计算分片总数
        file_parts[file_name] = {"received_parts": set(), "total_parts": total_parts}

    file_parts[file_name]["received_parts"].add(part_number)

    # 检查是否已接收所有分片
    if len(file_parts[file_name]["received_parts"]) == file_parts[file_name]["total_parts"]:
        await merge_file(file_name, file_parts[file_name]["total_parts"])
        del file_parts[file_name]  # 清理记录



    return web.Response(text=f"Part {part_number} received")


async def merge_file(file_name, total_parts):
    print(f'total_parts={total_parts}')
    """合并所有分片文件"""
    file_path = os.path.join(f'{args.temp_dir}', file_name)
    with open(file_path, "wb") as output_file:
        for part_number in range(1, total_parts + 1):
            part_path = f"{file_path}.part{part_number}"
            with open(part_path, "rb") as part_file:
                output_file.write(part_file.read())
            os.remove(part_path)  # 合并后删除分片文件
    print(f"{file_name} merged successfully")


app = web.Application(client_max_size=20 * 1024 * 1024)
app.add_routes([web.post('/recognition_chunk', api_recognition_chunk)])

if len(args.certfile) > 0:
    ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)

    # Generate with Lets Encrypt, copied to this location, chown to current user and 400 permissions
    ssl_cert = args.certfile
    ssl_key = args.keyfile

    ssl_context.load_cert_chain(ssl_cert, keyfile=ssl_key)
    web.run_app(app, host=args.host, port=args.port, ssl_context=ssl_context)
else:
    web.run_app(app, host=args.host, port=args.port)

